import torch
from torch import nn
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from torch.nn.modules.utils import _pair

# import alphaction._custom_cuda_ext as _C
from torchvision.ops import roi_align


# class _ROIAlign3d(Function):
#     @staticmethod
#     def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio):
#         ctx.save_for_backward(roi)
#         ctx.output_size = _pair(output_size)
#         ctx.spatial_scale = spatial_scale
#         ctx.sampling_ratio = sampling_ratio
#         ctx.input_shape = input.size()
#         output = _C.roi_align_3d_forward(
#             input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio
#         )
#         return output
#
#     @staticmethod
#     @once_differentiable
#     def backward(ctx, grad_output):
#         rois, = ctx.saved_tensors
#         output_size = ctx.output_size
#         spatial_scale = ctx.spatial_scale
#         sampling_ratio = ctx.sampling_ratio
#         bs, ch, l, h, w = ctx.input_shape
#         grad_input = _C.roi_align_3d_backward(
#             grad_output,
#             rois,
#             spatial_scale,
#             output_size[0],
#             output_size[1],
#             bs,
#             ch,
#             l,
#             h,
#             w,
#             sampling_ratio,
#         )
#         return grad_input, None, None, None, None
#
#
# roi_align_3d = _ROIAlign3d.apply


class ROIAlign3d(nn.Module):
    def __init__(self, output_size, spatial_scale, sampling_ratio):
        super(ROIAlign3d, self).__init__()
        self.output_size = output_size
        self.spatial_scale = spatial_scale
        self.sampling_ratio = sampling_ratio

    def forward(self, input, rois):
        # return roi_align_3d(
        #     input, rois, self.output_size, self.spatial_scale, self.sampling_ratio
        # )
        # aia提供的algin不支持半精度，先使用torchvision
        return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio)

    def __repr__(self):
        tmpstr = self.__class__.__name__ + "("
        tmpstr += "output_size=" + str(self.output_size)
        tmpstr += ", spatial_scale=" + str(self.spatial_scale)
        tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)
        tmpstr += ")"
        return tmpstr
